"""
tool工具的基本使用
"""
import asyncio
from typing import Any, Type
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.tools import tool, StructuredTool, BaseTool, ToolException
from pydantic import Field, BaseModel

from models import get_ds_model_client


# 注解 @tool
@tool
def multiply(a: int, b: int) -> int:
    """ multiply two numbers """
    return a * b


# print(multiply.name)
# print(multiply.description)
# print(multiply.args)
# print("-" * 80)

# 结构化工具 StructuredTool.from_function
def multiply_tool(a: int, b: int) -> int:
    """ multiply two numbers """
    return a * b


calculator = StructuredTool.from_function(func=multiply_tool)
# print(calculator.invoke({"a": 2, "b": 3}))
# print("-" * 80)

# 从可运行对象创建工具
def runnable_obj_tool():
    prompt = ChatPromptTemplate.from_messages(
        [("human", "Hello. Please respond in the style of {answer_style}.")]
    )
    llm = get_ds_model_client()
    parser = StrOutputParser()
    chain = prompt | llm | parser
    as_tool = chain.as_tool(
        name="Style responder", description="Description of when to use tool."
    )
    print(as_tool.args)
    # print(as_tool.invoke({"answer_style": "formal"}))
    print("-" * 80)

# runnable_obj_tool()

# 自定义工具 继承子类化 BaseTool
class CustomInput(BaseModel):
    a: int = Field(description="param a")
    b: int = Field(description="param b")
class CustomTool(BaseTool):
    name: str = "customTool"
    description: str = "sync multiply two numbers, async sum two numbers"
    args_schema: Type[BaseModel] = CustomInput

    def _run(self, a: int, b: int) -> Any:
        """Use the tool."""
        return a * b

    async def _arun(self, a: int, b: int) -> Any:
        return a + b

custom_tool = CustomTool()
# 同步
# print(custom_tool.invoke({"a": 2, "b": 3}))
# 异步
# print(asyncio.run(custom_tool.ainvoke({"a": 2, "b": 3})))
# print("-" * 80)


# 处理工具错误
def get_weather(city: str) -> int:
    """Get weather for the given city."""
    raise ToolException(f"Error: There is no city by the name of {city}.")
# 使用函数处理错误
def _handle_error(error: ToolException) -> str:
    return f"The following errors occurred during tool execution: `{error.args[0]}`"

# 如果不加handle_tool_error=True，会打印Traceback栈异常信息
get_weather_tool = StructuredTool.from_function(func=get_weather, handle_tool_error=True)
# 可以将 handle_tool_error 设置为一个始终返回的字符串。
# get_weather_tool = StructuredTool.from_function(func=get_weather, handle_tool_error="There is no such city, but it's probably above 0K there!")
# 使用函数处理错误
# get_weather_tool = StructuredTool.from_function(func=get_weather, handle_tool_error=_handle_error)
print(get_weather_tool.invoke({"city": "New York"}))


